import torch
from model import ShuffleNetV2
from datasets import MyDataset
# from utils.loss import bounding_box_loss

lr = 1e-5
train_dataset = MyDataset()
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
print('Dataset loading complete.')
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = ShuffleNetV2()
optimizer = torch.optim.Adam(model.parameters(), lr)
loss_function = torch.nn.CrossEntropyLoss()
weights_path = 'shufflenet.pth'
model.load_state_dict(torch.load(weights_path))
model.to(device)
print('Network loading complete.')
train_loss = 0
epoch = 0
while True:
    for step, data in enumerate(train_dataloader, start=0):
        images, labels = data
        preds_train = model(images.to(device))
        train_loss = loss_function(preds_train, labels.to(device))
        optimizer.zero_grad()
        train_loss.backward()
        optimizer.step()
        print(f'\rloss={train_loss.item()}', end='')
    print(f'\nEpoch{epoch} finished')
    epoch += 1
